import cv2
from paddleocr import PaddleOCR


def getPlate(image, fileName):
    rawImage = image.copy()
    # cv2.imshow("original", image)  # 测试语句，观察原始图像
    # 去噪处理
    image = cv2.GaussianBlur(image, (3, 3), 0)
    # cv2.imshow("GaussianBlur", image)  # 测试语句，查看滤波结果（去噪）

    filename = "out/%s/1.jpg" % (fileName.split(".")[0])
    cv2.imwrite(filename, image)

    # cv2.imwrite()函数返回一个布尔值，表示保存是否成功
    # 色彩空间转换（RGB-->GRAY)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    # cv2.imshow("gray", image)  # 测试语句，查看灰度图像

    filename = "out/%s/2.jpg" % (fileName.split(".")[0])
    cv2.imwrite(filename, image)

    # Sobel算子（X方向边缘梯度）
    SobelX = cv2.Sobel(image, cv2.CV_16S, 1, 0)
    absX = cv2.convertScaleAbs(SobelX)  # 映射到[0.255]内
    image = absX
    # cv2.imshow("soblex", image)  # 测试语句，图像边缘

    filename = "out/%s/3.jpg" % (fileName.split(".")[0])
    cv2.imwrite(filename, image)

    # 阈值处理
    ret, image = cv2.threshold(image, 0, 255, cv2.THRESH_OTSU)
    # cv2.imshow("imageThreshold", image)  # 测试语句，查看处理结果

    filename = "out/%s/4.jpg" % (fileName.split(".")[0])
    cv2.imwrite(filename, image)

    # 闭运算：先膨胀后腐蚀，车牌各个字符是分散的，让车牌构成一体
    kernelX = cv2.getStructuringElement(cv2.MORPH_RECT, (17, 5))
    image = cv2.morphologyEx(image, cv2.MORPH_CLOSE, kernelX)
    # cv2.imshow("imageCLOSE", image)  # 测试语句，查看处理结果

    filename = "out/%s/5.jpg" % (fileName.split(".")[0])
    cv2.imwrite(filename, image)

    # 开运算：先腐蚀后膨胀，去除噪声
    kernelY = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 19))
    image = cv2.morphologyEx(image, cv2.MORPH_OPEN, kernelY)
    cv2.imshow("imageOPEN", image)

    filename = "out/%s/6.jpg" % (fileName.split(".")[0])
    cv2.imwrite(filename, image)

    # 中值滤波：去除噪声
    image = cv2.medianBlur(image, 15)
    cv2.imshow("imagemedianBlur", image)  # 测试语句，查看处理结果

    filename = "out/%s/7.jpg" % (fileName.split(".")[0])

    cv2.imwrite(filename, image)

    # 查找轮廓
    contours, w1 = cv2.findContours(image, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    # 测试语句，查看轮廓
    image = cv2.drawContours(rawImage.copy(), contours, -1, (0, 0, 255), 3)
    cv2.imshow('imagecc', image)

    filename = "out/%s/8.jpg" % (fileName.split(".")[0])

    cv2.imwrite(filename, image)

    # 逐个遍历，将宽度>3倍高度的轮廓确定为车牌
    for item in contours:
        rect = cv2.boundingRect(item)
        x = rect[0]
        y = rect[1]
        weight = rect[2]
        height = rect[3]
        if weight > (height * 3):
            plate = rawImage[y:y + height, x:x + weight]

    # 使用CPU预加载，不用GPU
    cv2.imshow('plate', plate)  # 测试语句：查看提取车牌

    filename = "out/%s/9.jpg" % (fileName.split(".")[0])

    cv2.imwrite(filename, plate)
    ocr = PaddleOCR(use_angle_cls=True, use_gpu=False, ocr_version='PP-OCRv3')
    text = ocr.ocr(plate, cls=True)
    # 将txt的结果保存在文件中
    print(text)
    with open('out/%s/10.txt' % fileName.split('.')[0], 'w', encoding='utf-8') as f:
        for t in text:
            f.write(t[0][1][0])
            f.write('\n')
